#include "nvds_pipeline_builder.h"
#include "nvds_pipeline.h"
#include <memory>
#include <spdlog/spdlog.h>
#include <stdexcept>
#include <cmath>
#include <nvds_yml_parser.h>

#define CHECK_PARSER_ERROR(parse_expr)                                       \
    do {                                                                     \
        if (NVDS_YAML_PARSER_SUCCESS != parse_expr)                          \
            throw std::runtime_error("Error in parsing configuration file"); \
    } while (0)

GstElement *GstPipelineBuilder::buildPipelineWithDisplay(const char         *pipeline_name,
                                                         GstElement         *streammux,
                                                         std::string         nvds_config,
                                                         uint32_t            num_sources,
                                                         GstPadProbeCallback nvinfer_cb,
                                                         void               *nvinfer_args,
                                                         GstPadProbeCallback tailer_cb,
                                                         void               *tailer_args)
{
    auto pgie      = std::make_shared<NVDSElement>("nvinfer", "primary-nvinference-engine");
    auto queue1    = std::make_shared<NVDSElement>("queue", "queue1");
    auto queue2    = std::make_shared<NVDSElement>("queue", "queue2");
    auto queue3    = std::make_shared<NVDSElement>("queue", "queue3");
    auto queue4    = std::make_shared<NVDSElement>("queue", "queue4");
    auto queue5    = std::make_shared<NVDSElement>("queue", "queue5");
    auto tiler     = std::make_shared<NVDSElement>("nvmultistreamtiler", "nvtiler");
    auto nvvidconv = std::make_shared<NVDSElement>("nvvideoconvert", "nvvideo-converter");
    auto nvosd     = std::make_shared<NVDSElement>("nvdsosd", "nv-onscreendisplay");

#ifdef __aarch64__
    auto sink = std::make_shared<NVDSElement>("nv3dsink", "nv3d-sink");
#else
    auto sink = std::make_shared<NVDSElement>("nveglglessink", "nvvideo-renderer");
#endif

    // 绑定回调函数
    pgie->addSrcProbe(nvinfer_cb, nvinfer_args);
    sink->addSinkProbe(tailer_cb, tailer_args);

    GstElement *pipeline = gst_pipeline_new(pipeline_name);
    if (!pipeline)
        throw std::runtime_error("Create gstreamer pipeline failed!");
    gst_bin_add(GST_BIN(pipeline), streammux);

    // 设置配置参数
    CHECK_PARSER_ERROR(nvds_parse_streammux(streammux, (gchar *)nvds_config.c_str(), "streammux"));
    CHECK_PARSER_ERROR(nvds_parse_gie(pgie->getElement(), (gchar *)nvds_config.c_str(), "primary-gie"));

    // 设置可视化样式
    guint viz_rows    = (guint)std::sqrt(num_sources);
    guint viz_columns = (guint)std::ceil(1.0 * num_sources / viz_rows);
    g_object_set(G_OBJECT(tiler->getElement()), "rows", viz_rows, "columns", viz_columns, NULL);

    // 修正 batch size
    guint pgie_batch_size;
    g_object_get(G_OBJECT(pgie->getElement()), "batch-size", &pgie_batch_size, NULL);
    if (pgie_batch_size != num_sources) {
        spdlog::warn("WARNING: Overriding infer-config batch-size ({}) with number of sources ({})", pgie_batch_size, num_sources);
        g_object_set(G_OBJECT(pgie->getElement()), "batch-size", num_sources, NULL);
    }

    gst_bin_add_many(GST_BIN(pipeline),
                     queue1->getElement(),
                     pgie->getElement(),
                     queue2->getElement(),
                     tiler->getElement(),
                     queue3->getElement(),
                     nvvidconv->getElement(),
                     queue4->getElement(),
                     nvosd->getElement(),
                     queue5->getElement(),
                     sink->getElement(),
                     NULL);

    if (!gst_element_link_many(streammux,
                               queue1->getElement(),
                               pgie->getElement(),
                               queue2->getElement(),
                               tiler->getElement(),
                               queue3->getElement(),
                               nvvidconv->getElement(),
                               queue4->getElement(),
                               nvosd->getElement(),
                               queue5->getElement(),
                               sink->getElement(),
                               NULL)) {
        throw std::runtime_error("Elements could not be linked. Exiting.");
    }
    return pipeline;
}

GstElement *GstPipelineBuilder::buildPipeline(const char         *pipeline_name,
                                              GstElement         *streammux,
                                              std::string         nvds_config,
                                              uint32_t            num_sources,
                                              GstPadProbeCallback nvinfer_cb,
                                              void               *nvinfer_args,
                                              GstPadProbeCallback tailer_cb,
                                              void               *tailer_args)
{
    auto pgie   = std::make_shared<NVDSElement>("nvinfer", "primary-nvinference-engine");
    auto queue1 = std::make_shared<NVDSElement>("queue", "queue1");
    auto queue2 = std::make_shared<NVDSElement>("queue", "queue2");
    auto sink   = std::make_shared<NVDSElement>("fakesink", "nvvideo-renderer");

    // 绑定回调函数
    pgie->addSrcProbe(nvinfer_cb, nvinfer_args);
    sink->addSinkProbe(tailer_cb, tailer_args);

    GstElement *pipeline = gst_pipeline_new(pipeline_name);
    if (!pipeline)
        throw std::runtime_error("Create gstreamer pipeline failed!");
    gst_bin_add(GST_BIN(pipeline), streammux);

    // 设置配置参数
    CHECK_PARSER_ERROR(nvds_parse_streammux(streammux, (gchar *)nvds_config.c_str(), "streammux"));
    CHECK_PARSER_ERROR(nvds_parse_gie(pgie->getElement(), (gchar *)nvds_config.c_str(), "primary-gie"));

    // 修正 batch size
    guint pgie_batch_size;
    g_object_get(G_OBJECT(pgie->getElement()), "batch-size", &pgie_batch_size, NULL);
    if (pgie_batch_size != num_sources) {
        spdlog::warn("WARNING: Overriding infer-config batch-size ({}) with number of sources ({})", pgie_batch_size, num_sources);
        g_object_set(G_OBJECT(pgie->getElement()), "batch-size", num_sources, NULL);
    }

    gst_bin_add_many(GST_BIN(pipeline), queue1->getElement(), pgie->getElement(), queue2->getElement(), sink->getElement(), NULL);
    if (!gst_element_link_many(streammux, queue1->getElement(), pgie->getElement(), queue2->getElement(), sink->getElement(), NULL)) {
        throw std::runtime_error("Elements could not be linked. Exiting.");
    }
    return pipeline;
}